#include "UtilityPCH.h"
#include "UT_BSDF.h"
#include "UT_MathHelper.h"
#include "UT_SampleHelper.h"

using namespace Utility;
using namespace Memory;


//----------------------------------------------------------------------
//UT_BSDFSampleOffset
//----------------------------------------------------------------------
UT_BSDFSampleOffsets::UT_BSDFSampleOffsets(int count, UT_Sample* sample)
{
	nSamples = count;
	componentOffset = sample->Add1D(count);
	dirOffset = sample->Add2D(count);
}
//----------------------------------------------------------------------
//UT_BSDFSample
//----------------------------------------------------------------------
UT_BSDFSample::UT_BSDFSample(const UT_Sample* sample, const UT_BSDFSampleOffsets& offsets, uint32_t num)
{
	UT_ASSERT(num < sample->m_k2DNumVector.GetAt(offsets.dirOffset));
	UT_ASSERT(num < sample->m_k1DNumVector.GetAt(offsets.componentOffset));
	
	uDir[0]			= sample->m_pp2DData[offsets.dirOffset][2*num];
	uDir[1]			= sample->m_pp2DData[offsets.dirOffset][2*num+1];
	uComponent		= sample->m_pp1DData[offsets.componentOffset][num];
	
	UT_ASSERT(uDir[0] >= 0.f && uDir[0] < 1.f);
	UT_ASSERT(uDir[1] >= 0.f && uDir[1] < 1.f);
	UT_ASSERT(uComponent >= 0.f && uComponent < 1.f);
}
//----------------------------------------------------------------------
//UT_BSDF
//----------------------------------------------------------------------
UT_BSDF::UT_BSDF(const UT_DifferentialGeometry& rkDgs, const UT_Normal3Float& rkNgeom, float fEta /* = 1.f */)
	: m_kDgShading(rkDgs)
	, m_kNG(rkNgeom)
	, m_fEta(fEta)
{
	m_kNG = rkNgeom;
	m_kNS = rkDgs.nn;
	m_kSN = UT_MathHelper::Normalize(rkDgs.dpdu);
	m_kTN = UT_MathHelper::Cross(m_kNS, m_kSN);
	
	m_iBxDFCount = 0;
}
//----------------------------------------------------------------------
void UT_BSDF::AddBxDF(UT_BxDFBase* pkBxDF)
{
	UT_ASSERT(m_iBxDFCount < MAX_BxDFS_IN_BSDF);
	m_apBxDFs[m_iBxDFCount++]=pkBxDF;
}
//----------------------------------------------------------------------
int32_t UT_BSDF::GetBxDFCount() const
{
	return m_iBxDFCount;
}
//----------------------------------------------------------------------
int32_t UT_BSDF::GetBxDFCount(UT_BxDFBase::BxDFType eFlags) const
{
	int32_t iCount = 0;
	for (int32_t i = 0; i < m_iBxDFCount; ++i)
	{
		if (m_apBxDFs[i]->MatchesFlags(eFlags))
		{
			++iCount;
		}
	}
	return iCount;
}
//----------------------------------------------------------------------
UT_Vec3Float UT_BSDF::WorldToLocal(const UT_Vec3Float& v) const
{
	return UT_Vec3Float(UT_MathHelper::Dot(v, m_kSN), UT_MathHelper::Dot(v, m_kTN), UT_MathHelper::Dot(v, m_kNS));
}
//----------------------------------------------------------------------
UT_Vec3Float UT_BSDF::LocalToWorld(const UT_Vec3Float& v) const
{
	return UT_Vec3Float(m_kSN.x * v.x + m_kTN.x * v.y + m_kNS.x * v.z,
		m_kSN.y * v.x + m_kTN.y * v.y + m_kNS.y * v.z,
		m_kSN.z * v.x + m_kTN.z * v.y + m_kNS.z * v.z);
}
//----------------------------------------------------------------------
UT_SpectrumUsed UT_BSDF::f(const UT_Vec3Float& woW, const UT_Vec3Float& wiW, UT_BxDFBase::BxDFType eFlags /* = UT_BxDFBase::BSDF_ALL */) const
{
	UT_Vec3Float wi = WorldToLocal(wiW);
	UT_Vec3Float wo = WorldToLocal(woW);
	
	if (UT_MathHelper::Dot(wiW, m_kNG) * UT_MathHelper::Dot(woW, m_kNG) > 0) // ignore BTDFs
		eFlags = UT_BxDFBase::BxDFType(eFlags & ~UT_BxDFBase::BSDF_TRANSMISSION);
	else // ignore BRDFs
		eFlags = UT_BxDFBase::BxDFType(eFlags & ~UT_BxDFBase::BSDF_REFLECTION);

	UT_SpectrumUsed f = 0.;
	for (int i = 0; i < m_iBxDFCount; ++i)
		if (m_apBxDFs[i]->MatchesFlags(eFlags))
			f += m_apBxDFs[i]->f(wo, wi);

	return f;
}
//----------------------------------------------------------------------
UT_SpectrumUsed UT_BSDF::Sample_f(const UT_Vec3Float& woW, UT_Vec3Float* wiW, const UT_BSDFSample& rkBSDFSample, float* fPdf, UT_BxDFBase::BxDFType eFlags /* = UT_BxDFBase::BSDF_ALL */, UT_BxDFBase::BxDFType* peSampledType /* = NULL */) const
{
	// Choose which _BxDF_ to sample
	int32_t iMatchingComps = GetBxDFCount(eFlags);
	
	if (iMatchingComps == 0)
	{
		*fPdf = 0.f;
		if (peSampledType)
			*peSampledType = UT_BxDFBase::BxDFType(0);

		return UT_SpectrumUsed(0.f);
	}

	int32_t iWhichBxDF  = UT_MathHelper::Min(UT_MathHelper::Floor2Int(rkBSDFSample.uComponent * iMatchingComps),iMatchingComps-1);
	UT_BxDFBase* pkBxDF = NULL;
	int32_t iCount = iWhichBxDF;

	for (int i = 0; i < m_iBxDFCount; ++i)
	{
		if (m_apBxDFs[i]->MatchesFlags(eFlags) && iCount-- == 0)
		{
			pkBxDF = m_apBxDFs[i];
			break;
		}
	}
	UT_ASSERT(pkBxDF);
	
	// Sample chosen _BxDF_
	UT_Vec3Float wo = WorldToLocal(woW);
	UT_Vec3Float wi;
	*fPdf = 0.f;
	
	UT_SpectrumUsed f = pkBxDF->Sample_f(wo, &wi, rkBSDFSample.uDir[0],	rkBSDFSample.uDir[1], fPdf);
	
	if (*fPdf == 0.f)
	{
		if (peSampledType)
			*peSampledType = UT_BxDFBase::BxDFType(0);
		return 0.f;
	}
	if (peSampledType)
		*peSampledType = pkBxDF->m_eBxDFType;

	*wiW = LocalToWorld(wi);

	// Compute overall PDF with all matching _BxDF_s
	if (!(pkBxDF->m_eBxDFType & UT_BxDFBase::BSDF_SPECULAR) && iMatchingComps > 1)
	{
		for (int i = 0; i < m_iBxDFCount; ++i)
		{
			if (m_apBxDFs[i] != pkBxDF && m_apBxDFs[i]->MatchesFlags(eFlags))
			{
				*fPdf += m_apBxDFs[i]->Pdf(wo, wi);
			}
		}
	}

	if (iMatchingComps > 1)
		*fPdf /= iMatchingComps;

	// Compute value of BSDF for sampled direction
	if (!(pkBxDF->m_eBxDFType & UT_BxDFBase::BSDF_SPECULAR))
	{
		f = UT_BSDF::f(woW,*wiW,eFlags);
	}

	return f;
}
//----------------------------------------------------------------------
UT_SpectrumUsed UT_BSDF::rho(UT_RandomNumGenerator& rng, UT_BxDFBase::BxDFType eFlags /* = UT_BxDFBase::BSDF_ALL */, int32_t iSqrtSamples /* = 6 */) const
{
	int32_t nSamples = iSqrtSamples * iSqrtSamples;

	float *s1 = (float*)(MM_STAND_ALLOC(sizeof(float) * 2 * nSamples));
	UT_SampleHelper::SampleStratified2D(s1, iSqrtSamples, iSqrtSamples, rng);
	
	float *s2 = (float*)(MM_STAND_ALLOC(sizeof(float) * 2 * nSamples));
	UT_SampleHelper::SampleStratified2D(s2, iSqrtSamples, iSqrtSamples, rng);

	UT_SpectrumUsed ret(0.);
	for (int32_t i = 0; i < m_iBxDFCount; ++i)
		if (m_apBxDFs[i]->MatchesFlags(eFlags))
			ret += m_apBxDFs[i]->rho(nSamples, s1, s2);
	
	return ret;
}
//----------------------------------------------------------------------
UT_SpectrumUsed UT_BSDF::rho(const UT_Vec3Float& wo, UT_RandomNumGenerator& rng, UT_BxDFBase::BxDFType eFlags /* = UT_BxDFBase::BSDF_ALL */, int32_t iSqrtSamples /* = 6 */) const
{
	int32_t nSamples = iSqrtSamples * iSqrtSamples;
	
	float *s1 = (float*)(MM_STAND_ALLOC(sizeof(float) * 2 * nSamples));
	UT_SampleHelper::SampleStratified2D(s1, iSqrtSamples, iSqrtSamples, rng);

	UT_SpectrumUsed ret(0.);
	for (int32_t i = 0; i < m_iBxDFCount; ++i)
		if (m_apBxDFs[i]->MatchesFlags(eFlags))
			ret += m_apBxDFs[i]->rho(wo, nSamples, s1);
	
	return ret;
}
//----------------------------------------------------------------------
float UT_BSDF::Pdf(const UT_Vec3Float& woW, const UT_Vec3Float& wiW, UT_BxDFBase::BxDFType eFlags /* = UT_BxDFBase::BSDF_ALL */) const
{
	if (m_iBxDFCount == 0.)
		return 0.;

	UT_Vec3Float wo = WorldToLocal(woW);
	UT_Vec3Float wi = WorldToLocal(wiW);
	
	float	pdf = 0.f;
	int32_t matchingComps = 0;
	for (int32_t i = 0; i < m_iBxDFCount; ++i)
	{
		if (m_apBxDFs[i]->MatchesFlags(eFlags))
		{
			++matchingComps;
			pdf += m_apBxDFs[i]->Pdf(wo, wi);
		}
	}
	
	float v = matchingComps > 0 ? pdf / matchingComps : 0.f;
	return v;
}
//----------------------------------------------------------------------